use std::sync::{Arc, Mutex};

use async_trait::async_trait;
use pingora::{http::ResponseHeader, prelude::*};

pub struct LBService(pub Arc<LoadBalancer<RoundRobin>>, pub Mutex<usize>);

pub struct ReqCtx {
    beta_user: bool,
}

#[async_trait]
impl ProxyHttp for LBService {
    // 定义CTX类型
    type CTX = ReqCtx;

    fn new_ctx(&self) -> Self::CTX {
        ReqCtx { beta_user: false }
    }

    async fn upstream_peer(
        &self,
        session: &mut Session,
        ctx: &mut Self::CTX,
    ) -> Result<Box<HttpPeer>> {
        // 记录请求数
        let mut cnt = __self.1.lock().unwrap();
        *cnt += 1;

        if ctx.beta_user {
            println!("beta user request.")
        }

        let upstream = match self.0.select(b"", 256) {
            Some(upstream) => upstream,
            None => return Err(Error::new(ErrorType::new("no upstream"))),
        };

        let uri = &session.req_header().uri;
        println!(
            "select upstream peer is: {upstream:?} for {uri}, cnt:{}",
            cnt
        );

        let peer: Box<HttpPeer> = Box::new(HttpPeer::new(upstream, false, "".to_string()));
        Ok(peer)
    }

    async fn request_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<bool> {
        // 判断请求头
        ctx.beta_user = session.req_header().headers.get("beta-flag").is_some();
        Ok(false)
    }

    async fn response_filter(
        &self,
        session: &mut Session,
        upstream_response: &mut ResponseHeader,
        ctx: &mut Self::CTX,
    ) -> Result<()> {
        // replace existing header if any
        upstream_response
            .insert_header("Server", "MyGateway")
            .unwrap();
        // because we don't support h3
        upstream_response.remove_header("alt-svc");
        Ok(())
    }
}
